package com.embeddings.embed;

import com.alibaba.fastjson.JSONObject;
import okhttp3.*;

import java.io.IOException;
import java.util.List;

public class ReRankClient {

    public static String reRank(String uri, String apiKey, List<String> textsList, String query) throws IOException {
        OkHttpClient client = new OkHttpClient();
        JSONObject requestBody = new JSONObject();
        String[] texts = textsList.toArray(new String[0]);
        requestBody.put("textList", texts);
        requestBody.put("query", query);
        // 创建请求
        MediaType mediaType = MediaType.parse("application/json; charset=utf-8");
        RequestBody body = RequestBody.Companion.create(requestBody.toJSONString(), mediaType);
        Request request = new Request.Builder()
                .url(uri)
                .addHeader("Authorization", "Bearer " + apiKey)
                .addHeader("Content-Type", "application/json")
                .post(body)
                .build();

        // 发送请求
        Response response = client.newCall(request).execute();
        if (!response.isSuccessful()) {
            throw new IOException("Unexpected code " + response);
        }
        return response.body().string();
    }
}